from setuptools import setup, Extension
import pybind11
import os
import platform

# 检测当前架构
arch_flags = []
if platform.machine() in ['x86_64', 'AMD64']:
    arch_flags = ['-march=native', '-DARCH_X86']
    if platform.system() == 'Darwin':
        arch_flags.append('-mmacosx-version-min=10.15')
elif platform.machine().startswith('arm') or platform.machine().startswith('aarch64'):
    arch_flags = ['-march=native', '-DARCH_ARM']

# 设置编译参数
extra_compile_args = ['-O3', '-fopenmp'] + arch_flags
extra_link_args = ['-fopenmp']

# 添加源文件
sources = [
    'wrapper.cpp',
    '../core/core.c',
    '../core/array.c',
    '../core/matrix.c',
    '../core/image.c',
    '../core/memory.c',
    '../core/arch/dispatch.c'
]

# 添加架构特定源文件
if 'ARCH_X86' in arch_flags:
    sources.extend([
        '../core/arch/x86/simd.c',
        '../core/arch/x86/avx.c'
    ])
elif 'ARCH_ARM' in arch_flags:
    sources.extend([
        '../core/arch/arm/neon.c',
        '../core/arch/arm/sve.c'
    ])

module = Extension(
    'core',
    sources=sources,
    include_dirs=[
        pybind11.get_include(),
        '../core',
        '../core/arch',
        '../core/arch/x86',
        '../core/arch/arm'
    ],
    language='c++',
    extra_compile_args=extra_compile_args,
    extra_link_args=extra_link_args
)

setup(
    name='core_module',
    version='1.0',
    description='High-performance computing core',
    ext_modules=[module],
)